#!/usr/bin/env python3
import os
import random
import argparse
import logging
import numpy as np
from tensorboardX import SummaryWriter

from libbots import data, model, utils

import torch
import torch.optim as optim
import torch.nn.functional as F

import ptan

SAVES_DIR = "saves"

BATCH_SIZE = 16
LEARNING_RATE = 1e-4 # 因为是从已经训练好的模型继续训练，所以学习率是从1e-4开始，较低
MAX_EPOCHES = 10000

log = logging.getLogger("train")


def run_test(test_data, net, end_token, device="cpu"):
    '''
    test_data: 测试对话数据， shape = (batch_size, [first 对话， [second 应答对话列表]])
    net：训练中的模型网络
    end_token： 结束token

    return 返回模型预测的结果分数
    '''
    bleu_sum = 0.0
    bleu_count = 0
    for p1, p2 in test_data:
        # 得到用于模型网路哦的输入序列
        input_seq = model.pack_input(p1, net.emb, device)
        # 编码得到隐藏层信息
        enc = net.encode(input_seq)
        # todo 为什么这里采用input_seq.data[0:1]作为应答的起始字符
        # 预测应答对话
        _, tokens = net.decode_chain_argmax(enc, input_seq.data[0:1], seq_len=data.MAX_TOKENS,
                                            stop_at_token=end_token)

        # 得到实际应答去除第一个字符的token id序列，因为应答对话列表首个字符是#BEG token
        ref_indices = [
            indices[1:]
            for indices in p2
        ]
        # 计算两个序列之间的BLEU分数 todo bleu是怎么计算分数的
        bleu_sum += utils.calc_bleu_many(tokens, ref_indices)
        bleu_count += 1
    return bleu_sum / bleu_count


if __name__ == "__main__":
    logging.basicConfig(format="%(asctime)-15s %(levelname)s %(message)s", level=logging.INFO)
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", required=True, help="Category to use for training. Empty string to train on full dataset")
    parser.add_argument("--cuda", action='store_true', default=False, help="Enable cuda")
    parser.add_argument("-n", "--name", required=True, help="Name of the run")
    parser.add_argument("-l", "--load", required=True, help="Load model and continue in RL mode")
    # 用于强化学习微调中，每个训练样本执行的解码迭代的数量，越高可以得到更加精确的PG，但是会对GPU 内存要求高
    parser.add_argument("--samples", type=int, default=4, help="Count of samples in prob mode")
    # 书上有说明为什么禁用P246
    parser.add_argument("--disable-skip", default=False, action='store_true', help="Disable skipping of samples with high argmax BLEU")
    args = parser.parse_args()
    device = torch.device("cuda" if args.cuda else "cpu")

    # 创建模型保存路径
    saves_path = os.path.join(SAVES_DIR, args.name)
    os.makedirs(saves_path, exist_ok=True)

    # 加载训练数据
    phrase_pairs, emb_dict = data.load_data(genre_filter=args.data)
    log.info("Obtained %d phrase pairs with %d uniq words", len(phrase_pairs), len(emb_dict))
    # 保存词典
    data.save_emb_dict(saves_path, emb_dict)
    end_token = emb_dict[data.END_TOKEN]
    train_data = data.encode_phrase_pairs(phrase_pairs, emb_dict)
    rand = np.random.RandomState(data.SHUFFLE_SEED)
    rand.shuffle(train_data)
    train_data, test_data = data.split_train_test(train_data)
    log.info("Training data converted, got %d samples", len(train_data))
    # 重新按照上一句对对话进行分组
    train_data = data.group_train_data(train_data)
    test_data = data.group_train_data(test_data)
    log.info("Train set has %d phrases, test %d", len(train_data), len(test_data))

    # 得到token id 对 单词 字典
    rev_emb_dict = {idx: word for word, idx in emb_dict.items()}

    # 构建短语模型，模型的结构和训练一致
    net = model.PhraseModel(emb_size=model.EMBEDDING_DIM, dict_size=len(emb_dict),
                            hid_size=model.HIDDEN_STATE_SIZE).to(device)
    log.info("Model: %s", net)

    writer = SummaryWriter(comment="-" + args.name)
    # 加载已经训练好的模型
    net.load_state_dict(torch.load(args.load))
    log.info("Model loaded from %s, continue training in RL mode...", args.load)

    # BEGIN token 创建一个开始tensor,查看beg字符的词嵌入张量，用于解码器启动解码
    beg_token = torch.LongTensor([emb_dict[data.BEGIN_TOKEN]]).to(device)

    # todo 标记TBMeanTracker，并将其添加到writer
    with ptan.common.utils.TBMeanTracker(writer, batch_size=100) as tb_tracker:
        optimiser = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)
        batch_idx = 0 # 标记当前训练到第几轮
        best_bleu = None
        for epoch in range(MAX_EPOCHES):
            random.shuffle(train_data)
            dial_shown = False # 仅打印每个epoch第一个对话结果

            total_samples = 0
            skipped_samples = 0
            bleus_argmax = [] # 由argmax生成的序列的BLEU分数
            bleus_sample = [] # 存储在sample采样中使用策略梯度预测到的应答序列bleu分数

            for batch in data.iterate_batches(train_data, BATCH_SIZE):
                batch_idx += 1
                optimiser.zero_grad()
                input_seq, input_batch, output_batch = model.pack_batch_no_out(batch, net.emb, device)
                # 和训练一样，依旧是传输PackedSequence类型上一句后，得到隐藏层编码
                enc = net.encode(input_seq)

                net_policies = [] # todo 作用 存储的预测到的应答原始概率分布列表，数量有samples个
                net_actions = [] # todo 作用 存储预测到的应答列表，数量有sample个
                net_advantages = [] # todo 作用 存储sample采样中使用策略梯度方法预测到的结果与在之前使用最大Q值预测到的结果之间的分数之差得到的优势值
                # 这里是将开始token转换为词向量
                beg_embedding = net.emb(beg_token)

                for idx, inp_idx in enumerate(input_batch):
                    # 遍历每一句输入对话
                    total_samples += 1
                    # 获取输入对话对应的所有的应答对话，组成对话列表，并全部都去除掉第一个字符
                    ref_indices = [
                        indices[1:]
                        for indices in output_batch[idx]
                    ]
                    # 获取当前对话的隐藏层编码
                    item_enc = net.get_encoded_item(enc, idx)
                    # 得到当前网络隐藏层能够预测到的应答对话token id概率分布以及应答对话token id序列，这里是将token id序列当作每一个状态下执行的action进行处理了吗？ todo
                    # 因为这里模型是已经训练好了，所以这里可以使用argmax得到模型的预测应答
                    r_argmax, actions = net.decode_chain_argmax(item_enc, beg_embedding, data.MAX_TOKENS,
                                                                stop_at_token=end_token)
                    # todo 还是这个问题，上一步中获取到的应答对话是否也是缺少第一个字符以及为什么要去除掉第一个字符
                    # 计算bleu评测分，这个分数用于后续强化学习中的策略基线
                    argmax_bleu = utils.calc_bleu_many(actions, ref_indices)
                    bleus_argmax.append(argmax_bleu)

                    # 如果开启了跳过最大bleu分数的开光，且当前估算的bleu已经达到了
                    # 设定的阈值，则跳过本次的训练轮次，重新进行下一轮
                    # 如果argmax的bleu分数太高，则不用于训练，因为分数太高了，该回答很完美，没有训练的必要
                    if not args.disable_skip and argmax_bleu > 0.99:
                        skipped_samples += 1
                        continue

                    # 显示当前训练的对话结果
                    if not dial_shown:
                        # 打印输入的对话
                        log.info("Input: %s", utils.untokenize(data.decode_words(inp_idx, rev_emb_dict)))
                        # 打印应答label的对话
                        ref_words = [utils.untokenize(data.decode_words(ref, rev_emb_dict)) for ref in ref_indices]
                        log.info("Refer: %s", " ~~|~~ ".join(ref_words))
                        # 打印预测到的应答对话
                        log.info("Argmax: %s, bleu=%.4f", utils.untokenize(data.decode_words(actions, rev_emb_dict)),
                                 argmax_bleu)

                    # samples是什么作用
                    # 进入强化学习训练，sample表示使用强化学习随机采样的次数，用于评测训练
                    for _ in range(args.samples):
                        r_sample, actions = net.decode_chain_sampling(item_enc, beg_embedding, data.MAX_TOKENS, stop_at_token=end_token)
                        # 计算预测到的结果的bleu分数
                        sample_bleu = utils.calc_bleu_many(actions, ref_indices)

                        if not dial_shown:
                            # 打印预测到的应答以及分数
                            log.info("Sample: %s, bleu=%.4f", utils.untokenize(data.decode_words(actions, rev_emb_dict)),
                                     sample_bleu)

                        net_policies.append(r_sample)
                        net_actions.extend(actions)
                        # todo 这里乘len(actions)的操作是做什么，是为了对齐actions中每一个字符吗
                        net_advantages.extend([sample_bleu - argmax_bleu] * len(actions))
                        bleus_sample.append(sample_bleu)
                    dial_shown = True

                # todo 这里的作用 仅仅只是保证net_policies有数据吗
                if not net_policies:
                    continue

                policies_v = torch.cat(net_policies) # sample循环中每一个序列解码的logits的列表
                actions_t = torch.LongTensor(net_actions).to(device) # sample循环中预测的应答序列列表
                adv_v = torch.FloatTensor(net_advantages).to(device) # sample循环中每次生成的应答序列优势值
                log_prob_v = F.log_softmax(policies_v, dim=1)
                # 这里会提高使用策略梯度应答优于原始模型的概率，反之降低低于原始模型的应答
                log_prob_actions_v = adv_v * log_prob_v[range(len(net_actions)), actions_t]
                # 添加符号求最小值
                loss_policy_v = -log_prob_actions_v.mean()

                loss_v = loss_policy_v
                loss_v.backward()
                optimiser.step()

                tb_tracker.track("advantage", adv_v, batch_idx)
                tb_tracker.track("loss_policy", loss_policy_v, batch_idx)
                tb_tracker.track("loss_total", loss_v, batch_idx)

            # 运行测试训练后的模型，得到测试机的bleu分数
            bleu_test = run_test(test_data, net, end_token, device)
            # 计算本轮训练最大argmax策略得到的bleu分数
            bleu = np.mean(bleus_argmax)
            # 记录到tensorboard
            writer.add_scalar("bleu_test", bleu_test, batch_idx)
            writer.add_scalar("bleu_argmax", bleu, batch_idx)
            writer.add_scalar("bleu_sample", np.mean(bleus_sample), batch_idx)
            writer.add_scalar("skipped_samples", skipped_samples / total_samples, batch_idx)
            writer.add_scalar("epoch", batch_idx, epoch)
            log.info("Epoch %d, test BLEU: %.3f", epoch, bleu_test)
            # 如果得到的bleu分数好于原先的，则保存更新模型
            if best_bleu is None or best_bleu < bleu_test:
                best_bleu = bleu_test
                log.info("Best bleu updated: %.4f", bleu_test)
                torch.save(net.state_dict(), os.path.join(saves_path, "bleu_%.3f_%02d.dat" % (bleu_test, epoch)))

            # 这里仅保存中间的训练过程
            if epoch % 10 == 0:
                torch.save(net.state_dict(), os.path.join(saves_path, "epoch_%03d_%.3f_%.3f.dat" % (epoch, bleu, bleu_test)))

    writer.close()
